#!/usr/bin/env python3
"""
A股全市场数据初始化脚本
用于获取所有A股股票的历史数据作为基础数据
"""

import sys
import os
import time
from pathlib import Path
from datetime import datetime, timedelta
from typing import List, Dict, Any
import json
import pandas as pd

# 添加src目录到Python路径
sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'src'))

from src.ashare_list import AShareListFetcher
from src.ashare_fetcher import AShareFetcher
from src.data_storage import DataStorage
from src.utils.config import get_settings, ensure_data_dirs
from src.utils.logger import get_logger

logger = get_logger("AShareInitialize")


class AShareInitializer:
    """A股数据初始化器"""

    def __init__(self):
        self.settings = get_settings()
        self.logger = logger
        self.list_fetcher = AShareListFetcher()
        self.data_fetcher = AShareFetcher()
        self.data_storage = DataStorage()
        self.progress_file = None
        self.initialize_config = self._load_initialize_config()

    def _load_initialize_config(self) -> Dict[str, Any]:
        """加载初始化配置"""
        return {
            'batch_size': 50,  # 每批处理的股票数量
            'delay_between_batches': 2.0,  # 批次间延迟（秒）
            'delay_between_stocks': 0.1,  # 股票间延迟（秒）
            'max_retries': 3,  # 最大重试次数
            'start_date': '2015-01-01',  # 开始日期
            'end_date': datetime.now().strftime('%Y-%m-%d'),  # 结束日期
            'interval': '1d',  # 数据间隔
            'skip_existing': True,  # 跳过已存在的数据
            'save_progress': True,  # 保存进度
            'filter_st_stocks': True,  # 过滤ST股票
            'filter_low_volume': True,  # 过滤低成交量股票
            'min_volume_threshold': 1000000,  # 最小成交量阈值
            'include_etf': False,  # 包含ETF（暂时关闭，因为baostock ETF支持不好）
            'exclude_cyb': True,  # 排除创业板
            'exclude_bse': True,  # 排除北交所
        }

    def get_filtered_stock_list(self) -> List[str]:
        """获取过滤后的股票列表"""
        try:
            self.logger.info("获取股票列表...")
            df = self.list_fetcher.get_all_stocks()

            if df.empty:
                return []

            # 过滤ST股票
            if self.initialize_config['filter_st_stocks']:
                df = df[~df['name'].str.contains('ST', na=False)]
                df = df[~df['name'].str.contains('PT', na=False)]
                self.logger.info(f"过滤ST股票后剩余: {len(df)} 只")

            # 过滤新三板（通常流动性较差）
            df = df[df['market_type'] != '新三板']

            # 过滤创业板和北交所（根据用户账户限制）
            df = df[df['market_type'] != '创业板']
            df = df[df['market_type'] != '北交所']

            # 获取ETF列表
            etf_list = []
            if self.initialize_config['include_etf']:
                etf_list = self._get_etf_list()
                self.logger.info(f"获取到 {len(etf_list)} 只ETF")

            final_symbols = df['symbol'].tolist() + etf_list
            self.logger.info(f"最终筛选股票数量: {len(final_symbols)} 只 (股票: {len(df)}, ETF: {len(etf_list)})")

            # 按市场类型分组
            market_groups = df.groupby('market_type')
            for market, group in market_groups:
                self.logger.info(f"  {market}: {len(group)} 只")

            if etf_list:
                self.logger.info(f"  ETF: {len(etf_list)} 只")

            return final_symbols

        except Exception as e:
            self.logger.error(f"获取过滤后股票列表失败: {e}")
            return []

    def _get_etf_list(self) -> List[str]:
        """获取ETF列表"""
        try:
            self.logger.info("获取ETF列表...")

            # 常见的ETF代码列表（沪市和深市）
            common_etfs = [
                # 沪市ETF
                '510050',  # 上证50ETF
                '510300',  # 华泰柏瑞300ETF
                '510500',  # 南方中证500ETF
                '510880',  # 上证红利ETF
                '511010',  # 国泰国债ETF
                '511220',  # 海富通上证可质押城投债ETF
                '511260',  # 国泰上证10年期国债ETF
                '511280',  # 富国上证10年期国债ETF
                '511380',  # 上证180ETF
                '512000',  # 华宝中证1000ETF
                '512100',  # 南方中证1000ETF
                '512170',  # 华泰柏瑞医疗ETF
                '512180',  # 华安上证180ETF
                '512200',  # 南方中证房地产ETF
                '512400',  # 南方中证有色金属ETF
                '512600',  # 南方中证银行ETF
                '512660',  # 国泰中证军工ETF
                '512670',  # 南方中证证券公司ETF
                '512690',  # 华夏中证500ETF
                '512800',  # 华泰柏瑞中证800ETF
                '512880',  # 华安中证银行ETF
                '512900',  # 华夏中证全指证券公司ETF
                '512980',  # 华夏中证全指证券公司ETF
                '515000',  # 华夏中证5GETF
                '515180',  # 易方达中证红利ETF
                '515450',  # 华泰柏瑞中证光伏产业ETF
                '515880',  # 招商中证白酒ETF
                '516160',  # 南方中证新能源ETF
                '516970',  # 华夏中证新能源汽车ETF
                '517080',  # 华夏中证新能源ETF
                '517180',  # 华泰柏瑞中证光伏产业ETF
                '518880',  # 华安黄金ETF
                '588000',  # 华夏上证科创板50ETF
                '588080',  # 易方达上证科创板50ETF
                '588200',  # 华泰柏瑞中证科创创业50ETF
                '588300',  # 华夏中证科创创业50ETF
                '588790',  # 博时中证科创创业50ETF
                '588800',  # 南方中证科创创业50ETF

                # 深市ETF
                '159901',  # 易方达深证100ETF
                '159915',  # 南方深证成指ETF
                '159919',  # 嘉实沪深300ETF
                '159920',  # 华夏恒生ETF
                '159928',  # 汇添富中证800ETF
                '159934',  # 易方达黄金ETF
                '159949',  # 南方中证500ETF
                '159908',  # 博时中证500ETF
                '159922',  # 嘉实中证400ETF
                '159925',  # 南方中证500ETF
                '159941',  # 广发中证全指可选消费ETF
                '159945',  # 广发中证全指工业ETF
                '159946',  # 广发中证全指原材料ETF
                '159947',  # 广发中证全指能源ETF
                '159948',  # 广发中证全指金融地产ETF
                '159950',  # 广发中证全指医药卫生ETF
                '159951',  # 广发中证全指信息技术ETF
                '159952',  # 广发中证全指电信业务ETF
                '159953',  # 广发中证全指公用事业ETF
                '159958',  # 南方中证银行ETF
                '159967',  # 广发中证军工ETF
                '159969',  # 华夏中证证券公司ETF
                '159975',  # 南方中证证券公司ETF
                '159976',  # 南方中证银行ETF
                '159977',  # 南方中证全指房地产ETF
                '159978',  # 南方中证全指有色金属ETF
                '159979',  # 南方中证全指消费ETF
                '159980',  # 南方中证全指医药卫生ETF
                '159981',  # 南方中证全指信息技术ETF
                '159982',  # 南方中证全指金融地产ETF
                '159983',  # 南方中证全指工业ETF
                '159984',  # 南方中证全指原材料ETF
                '159985',  # 南方中证全指能源ETF
                '159986',  # 南方中证全指公用事业ETF
                '159987',  # 南方中证全指电信业务ETF
                '159988',  # 南方中证全指可选消费ETF
                '159989',  # 南方中证全指主要消费ETF
                '159990',  # 华夏中证主要消费ETF
                '159991',  # 华夏中证可选消费ETF
                '159992',  # 华夏中证医药卫生ETF
                '159993',  # 华夏中证信息技术ETF
                '159994',  # 华夏中证金融地产ETF
                '159995',  # 华夏中证工业ETF
                '159996',  # 华夏中证原材料ETF
                '159997',  # 华夏中证能源ETF
                '159998',  # 华夏中证公用事业ETF
                '159999',  # 华夏中证电信业务ETF
            ]

            # 也可以尝试从baostock获取ETF列表
            try:
                if not self.list_fetcher.login():
                    return common_etfs

                # 查询ETF基本信息
                etf_rs = self.list_fetcher.bs.query_etf_basic()
                if etf_rs.error_code == '0':
                    etf_data = []
                    while (etf_rs.error_code == '0') & etf_rs.next():
                        etf_data.append(etf_rs.get_row_data())

                    if etf_data:
                        etf_df = pd.DataFrame(etf_data, columns=etf_rs.fields)
                        # 提取ETF代码
                        etf_codes = etf_df['code'].str.replace('sh.', '').str.replace('sz.', '').tolist()
                        self.logger.info(f"从baostock获取到 {len(etf_codes)} 只ETF")
                        return etf_codes

            except Exception as e:
                self.logger.warning(f"从baostock获取ETF列表失败，使用预设列表: {e}")

            return common_etfs

        except Exception as e:
            self.logger.error(f"获取ETF列表失败: {e}")
            return []

    def _get_stock_name_mapping(self) -> Dict[str, str]:
        """获取股票名称映射"""
        # 主要股票名称映射
        stock_names = {
            # 深圳主板
            '000001': '平安银行',
            '000002': '万科A',
            '000858': '五粮液',
            '000895': '格力电器',
            '002415': '海康威视',
            '002594': '比亚迪',
            '002714': '牧原股份',

            # 上海主板
            '600000': '浦发银行',
            '600036': '招商银行',
            '600519': '贵州茅台',
            '600887': '伊利股份',
            '600958': '华夏银行',
            '601318': '中国平安',
            '601398': '工商银行',
            '601857': '中国石油',
            '601988': '中国银行',

            # 科创板
            '688036': '传音控股',
            '688228': '安集科技',
            '688981': '中芯国际',

            # ETF
            '510050': '上证50ETF',
            '510300': '沪深300ETF',
            '510500': '中证500ETF',
            '518880': '黄金ETF',
            '588000': '科创50ETF',
            '159901': '深100ETF',
            '159949': '中证500ETF',
            '159915': '深成指ETF',
        }

        return stock_names

    def _validate_stock_exists(self, symbol: str) -> bool:
        """验证股票是否存在"""
        try:
            # 使用获取器直接验证股票代码
            return self.data_fetcher.validate_symbol(symbol)

        except Exception as e:
            self.logger.error(f"验证股票 {symbol} 存在性时发生错误: {e}")
            return False

    def check_existing_data(self, symbol: str) -> bool:
        """检查是否已存在数据"""
        try:
            info = self.data_storage.get_data_info(symbol, self.initialize_config['interval'])
            return info['file_exists'] and info['record_count'] > 0
        except Exception:
            return False

    def initialize_stock(self, symbol: str, retry_count: int = 0) -> bool:
        """初始化单个股票数据"""
        try:
            # 检查是否已存在数据
            if self.initialize_config['skip_existing'] and self.check_existing_data(symbol):
                self.logger.debug(f"跳过已存在的数据: {symbol}")
                return True

            # 验证股票是否存在
            if not self._validate_stock_exists(symbol):
                self.logger.warning(f"股票 {symbol} 不存在或已退市，跳过")
                return True  # 返回True但跳过保存

            # 判断是否为ETF
            is_etf = self._is_etf_symbol(symbol)
            stock_type = "ETF" if is_etf else "股票"
            self.logger.info(f"初始化{stock_type} {symbol} (重试 {retry_count + 1}/{self.initialize_config['max_retries']})")

            # 获取历史数据
            try:
                data = self.data_fetcher.fetch_history_data(
                    symbol,
                    self.initialize_config['interval'],
                    self.initialize_config['start_date'],
                    self.initialize_config['end_date']
                )

                if data.empty:
                    self.logger.info(f"股票 {symbol} 在指定时间范围内无交易数据，跳过")
                    return True  # 返回True但跳过保存，不算作失败
            except Exception as e:
                self.logger.info(f"股票 {symbol} 数据获取失败: {e}，跳过")
                return True  # 返回True但跳过保存，不算作失败

            # 对于股票，过滤低成交量（ETF通常不需要这个过滤）
            if not is_etf and self.initialize_config['filter_low_volume']:
                if 'Volume' in data.columns:
                    avg_volume = data['Volume'].mean()
                    if avg_volume < self.initialize_config['min_volume_threshold']:
                        self.logger.info(f"跳过低成交量股票 {symbol} (平均成交量: {avg_volume:.0f})")
                        return True  # 返回True但跳过保存

            # 获取股票名称
            stock_name = None
            try:
                # 首先尝试从映射表获取
                name_mapping = self._get_stock_name_mapping()
                stock_name = name_mapping.get(symbol)

                # 如果映射表中没有，尝试从baostock获取
                if not stock_name:
                    stock_info = self.list_fetcher.get_stock_info(symbol)
                    stock_name = stock_info.get('name', symbol)

                # 如果还是没有，使用股票代码
                if not stock_name:
                    stock_name = symbol

            except Exception:
                stock_name = symbol

            # 保存数据
            success = self.data_storage.save_data(symbol, data, self.initialize_config['interval'], stock_name)

            if success:
                self.logger.info(f"成功初始化{stock_type} {symbol}: {len(data)} 条记录")
                return True
            else:
                self.logger.error(f"保存 {symbol} 数据失败")
                return False

        except Exception as e:
            self.logger.error(f"初始化 {symbol} 失败 (重试 {retry_count + 1}): {e}")
            return False

    def _is_etf_symbol(self, symbol: str) -> bool:
        """判断是否为ETF代码"""
        if not symbol:
            return False

        # 沪市ETF以5开头，深市ETF以159开头
        if symbol.startswith('5') or symbol.startswith('159'):
            return True

        # 检查是否在预设ETF列表中
        etf_list = self._get_etf_list()
        return symbol in etf_list

    def save_progress(self, processed_symbols: List[str], failed_symbols: List[str],
                      skipped_symbols: List[str]):
        """保存进度信息"""
        try:
            progress_data = {
                'timestamp': datetime.now().isoformat(),
                'config': self.initialize_config,
                'processed': processed_symbols,
                'failed': failed_symbols,
                'skipped': skipped_symbols,
                'total_processed': len(processed_symbols),
                'total_failed': len(failed_symbols),
                'total_skipped': len(skipped_symbols)
            }

            progress_file = Path(self.settings['system']['data_directory']) / 'metadata' / 'initialization_progress.json'
            with open(progress_file, 'w', encoding='utf-8') as f:
                json.dump(progress_data, f, ensure_ascii=False, indent=2)

            self.logger.info(f"进度已保存到: {progress_file}")

        except Exception as e:
            self.logger.error(f"保存进度失败: {e}")

    def load_progress(self) -> Dict[str, Any]:
        """加载进度信息"""
        try:
            progress_file = Path(self.settings['system']['data_directory']) / 'metadata' / 'initialization_progress.json'
            if not progress_file.exists():
                return {}

            with open(progress_file, 'r', encoding='utf-8') as f:
                return json.load(f)

        except Exception as e:
            self.logger.error(f"加载进度失败: {e}")
            return {}

    def run_initialization(self, symbols: List[str] = None):
        """运行初始化"""
        try:
            start_time = datetime.now()
            self.logger.info("=" * 60)
            self.logger.info("开始A股全市场数据初始化")
            self.logger.info("=" * 60)

            # 获取股票列表
            if symbols is None:
                symbols = self.get_filtered_stock_list()

            if not symbols:
                self.logger.error("没有获取到股票列表")
                return False

            total_count = len(symbols)
            self.logger.info(f"总共需要初始化 {total_count} 只股票")

            # 加载进度
            progress = self.load_progress()
            processed_symbols = progress.get('processed', [])
            failed_symbols = progress.get('failed', [])
            skipped_symbols = progress.get('skipped', [])

            # 过滤已处理的股票
            remaining_symbols = [s for s in symbols if s not in processed_symbols]
            remaining_count = len(remaining_symbols)

            self.logger.info(f"已处理: {len(processed_symbols)}, 剩余: {remaining_count}")

            if remaining_count == 0:
                self.logger.info("所有股票已经处理完成")
                return True

            # 分批处理
            batch_size = self.initialize_config['batch_size']
            total_batches = (remaining_count + batch_size - 1) // batch_size

            for batch_idx in range(total_batches):
                batch_start = batch_idx * batch_size
                batch_end = min(batch_start + batch_size, remaining_count)
                batch_symbols = remaining_symbols[batch_start:batch_end]

                self.logger.info(f"处理批次 {batch_idx + 1}/{total_batches}: {len(batch_symbols)} 只股票")

                batch_processed = 0
                batch_failed = 0
                batch_skipped = 0

                for i, symbol in enumerate(batch_symbols):
                    self.logger.info(f"处理 {symbol} ({batch_start + i + 1}/{remaining_count})")

                    # 初始化股票数据
                    success = self.initialize_stock(symbol)

                    if success:
                        processed_symbols.append(symbol)
                        batch_processed += 1
                    else:
                        # 重试机制
                        retry_success = False
                        for retry in range(self.initialize_config['max_retries']):
                            time.sleep(2 ** retry)  # 指数退避
                            if self.initialize_stock(symbol, retry + 1):
                                processed_symbols.append(symbol)
                                batch_processed += 1
                                retry_success = True
                                break

                        if not retry_success:
                            failed_symbols.append(symbol)
                            batch_failed += 1

                    # 股票间延迟
                    time.sleep(self.initialize_config['delay_between_stocks'])

                # 批次间延迟
                if batch_idx < total_batches - 1:
                    time.sleep(self.initialize_config['delay_between_batches'])

                # 保存进度
                if self.initialize_config['save_progress']:
                    self.save_progress(processed_symbols, failed_symbols, skipped_symbols)

                # 批次统计
                self.logger.info(f"批次 {batch_idx + 1} 完成: 成功 {batch_processed}, 失败 {batch_failed}")

            # 最终统计
            end_time = datetime.now()
            duration = end_time - start_time

            self.logger.info("=" * 60)
            self.logger.info("初始化完成！")
            self.logger.info("=" * 60)
            self.logger.info(f"总耗时: {duration}")
            self.logger.info(f"总股票数: {total_count}")
            self.logger.info(f"成功处理: {len(processed_symbols)}")
            self.logger.info(f"失败数量: {len(failed_symbols)}")
            self.logger.info(f"跳过数量: {len(skipped_symbols)}")
            self.logger.info(f"成功率: {len(processed_symbols) / total_count * 100:.1f}%")

            if failed_symbols:
                self.logger.warning(f"失败的股票: {failed_symbols[:10]}...")  # 只显示前10个

            return len(failed_symbols) == 0

        except Exception as e:
            self.logger.error(f"初始化过程出现异常: {e}")
            return False

    def run_quick_initialization(self, sample_count: int = 50):
        """快速初始化（用于测试）"""
        try:
            symbols = self.get_filtered_stock_list()
            sample_symbols = symbols[:sample_count]

            self.logger.info(f"运行快速初始化，处理 {len(sample_symbols)} 只股票")

            return self.run_initialization(sample_symbols)

        except Exception as e:
            self.logger.error(f"快速初始化失败: {e}")
            return False

    def check_initialization_status(self) -> Dict[str, Any]:
        """检查初始化状态"""
        try:
            symbols = self.get_filtered_stock_list()
            total_count = len(symbols)

            initialized_count = 0
            failed_count = 0
            not_initialized = []

            for symbol in symbols:
                info = self.data_storage.get_data_info(symbol, '1d')
                if info['file_exists']:
                    if info['record_count'] > 0:
                        initialized_count += 1
                    else:
                        failed_count += 1
                else:
                    not_initialized.append(symbol)

            return {
                'total_stocks': total_count,
                'initialized': initialized_count,
                'failed': failed_count,
                'not_initialized': not_initialized,
                'completion_rate': initialized_count / total_count * 100 if total_count > 0 else 0,
                'sample_not_initialized': not_initialized[:10]  # 显示前10个未初始化的
            }

        except Exception as e:
            self.logger.error(f"检查初始化状态失败: {e}")
            return {}


def main():
    """主函数"""
    import argparse

    parser = argparse.ArgumentParser(description='A股全市场数据初始化')
    parser.add_argument('--mode', choices=['full', 'quick', 'status'], default='full',
                       help='运行模式: full=全量初始化, quick=快速初始化, status=检查状态')
    parser.add_argument('--sample', type=int, default=50,
                       help='快速初始化的样本数量')
    parser.add_argument('--symbols', nargs='+', help='指定要初始化的股票代码')

    args = parser.parse_args()

    initializer = AShareInitializer()

    try:
        if args.mode == 'status':
            status = initializer.check_initialization_status()
            print("\n=== A股初始化状态 ===")
            print(f"总股票数: {status['total_stocks']}")
            print(f"已初始化: {status['initialized']}")
            print(f"失败数量: {status['failed']}")
            print(f"未初始化: {len(status['not_initialized'])}")
            print(f"完成率: {status['completion_rate']:.1f}%")

            if status['sample_not_initialized']:
                print(f"\n未初始化的股票示例: {status['sample_not_initialized']}")

        elif args.mode == 'quick':
            print(f"运行快速初始化，样本数量: {args.sample}")
            success = initializer.run_quick_initialization(args.sample)
            sys.exit(0 if success else 1)

        elif args.mode == 'full':
            if args.symbols:
                print(f"初始化指定股票: {args.symbols}")
                success = initializer.run_initialization(args.symbols)
            else:
                print("运行全量初始化")
                success = initializer.run_initialization()
            sys.exit(0 if success else 1)

    except KeyboardInterrupt:
        print("\n用户中断初始化过程")
        sys.exit(1)
    except Exception as e:
        print(f"初始化过程出现错误: {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()